package ssau.knyazev.fft.modules.real;

import ssau.knyazev.modules.Generator;
import ssau.knyazev.modules.Printer;

public class RealB2FFTModule {

	
	/**
	 * 
	 * @param src Matrix of real numbers
	 * @return Matrix of complex numbers. src[2*i] - real and src[2*i+1] - img
	 */
	
	public static float[][] _fft(float[][] src) {
		int N = src.length;
		float[] wi_list = new float[4*src.length];
		float[] wj_list = new float[4*src.length];
		
		for (int i = 0; i < src.length; i++) {
			double power = -2*Math.PI*i/N;
			wi_list[4*i] = (float)Math.cos(power)/2;
			wi_list[4*i+1] = (float)Math.cos(power)/2;
			wi_list[4*i+2] = (float)Math.sin(power)/2;
			wi_list[4*i+3] = (float)Math.sin(power)/2;
			
			wj_list[4*i] = (float)Math.cos(power)/2;
			wj_list[4*i+1] = (float)Math.cos(power)/2;
			wj_list[4*i+2] = -(float)Math.sin(power)/2;
			wj_list[4*i+3] = (float)Math.sin(power)/2;
		}
		
		_fftU02U2(src, wi_list, wj_list);
		_fftU01U3(src, wi_list, wj_list);
		
		return src;
	}
	
	protected static void _fftU02U2(float[][] src, float[] wi_list, float[] wj_list){
		int size = 0;
		int N = src.length;
		for (int x = src.length; x > 1; x = x >> 1){
			size++;
		}
		for (int q = 0; q < size; q++) {
			int step = 1 << q;
			int step2 = step << 1;
			for (int down = 0; down < N - step; down++) {
				for (int right = 0; right < N - step; right++) {
					boolean flag = (down/step) % 2 + (right/step) % 2 > 0 ? false : true;
					if (flag) {
						int wiInd = (down % step2)*(N/step2)*4;
						int wjInd = (right % step2)*(N/step2)*4;
						float[] wi = {wi_list[wiInd], wi_list[wiInd+2]};
						float[] wj = {wj_list[wjInd], wj_list[wjInd+2]};
						
						int ru0Ind = 4*right;
						int ru2Ind = ru0Ind + 2;
						int rsu0Ind = 4*(right+step);
						int rsu2Ind = rsu0Ind+2;
						
						float a00u0 = src[down][ru0Ind];
						float a00u2 = src[down][ru2Ind];
						float a10u0 = src[down][ru0Ind];
						float a10u2 = src[down][ru2Ind];
						float a01u0 = src[down][ru0Ind];
						float a01u2 = src[down][ru2Ind];
						float a11u0 = src[down][ru0Ind];
						float a11u2 = src[down][ru2Ind];
						
						float[] t = new float[2];
						t[0] = 2*wi[0]*src[down+step][ru0Ind] - 2*wi[1]*src[down+step][ru2Ind];
						t[1] = 2*wi[0]*src[down+step][ru2Ind] + 2*wi[1]*src[down+step][ru0Ind];
						a00u0 += t[0];
						a00u2 += t[1];
						a10u0 -= t[0];
						a10u2 -= t[1];
						a01u0 += t[0];
						a01u2 += t[1];
						a11u0 -= t[0];
						a11u2 -= t[1];
										
						t[0] = 2*wj[0]*src[down][rsu0Ind] - 2*wj[1]*src[down][rsu2Ind];
						t[1] = 2*wj[0]*src[down][rsu2Ind] + 2*wj[1]*src[down][rsu0Ind];
						a00u0 += t[0];
						a00u2 += t[1];
						a10u0 += t[0];
						a10u2 += t[1];
						a01u0 -= t[0];
						a01u2 -= t[1];
						a11u0 -= t[0];
						a11u2 -= t[1];
						
						float[] t2 = new float[2];
						t[0] = 2*wj[0]*src[down+step][rsu0Ind] - 2*wj[1]*src[down+step][rsu2Ind];
						t[1] = 2*wj[0]*src[down+step][rsu2Ind] + 2*wj[1]*src[down+step][rsu0Ind];
						t2[0] = 2*wi[0]*t[0] - 2*wi[1]*t[1];
						t2[1] = 2*wi[0]*t[1] + 2*wi[1]*t[0];
						a00u0 += t2[0];
						a00u2 += t2[1];
						a10u0 -= t2[0];
						a10u2 -= t2[1];
						a01u0 -= t2[0];
						a01u2 -= t2[1];
						a11u0 += t2[0];
						a11u2 += t2[1];
						
						src[down][ru0Ind] = a00u0;
						src[down][ru2Ind] = a00u2;
						src[down+step][ru0Ind] = a10u0;
						src[down+step][ru2Ind] = a10u2;
						src[down][rsu0Ind] = a01u0;
						src[down][rsu2Ind] = a01u2;
						src[down+step][rsu0Ind] = a11u0;
						src[down+step][rsu2Ind] = a11u2;		
					}
				}
			}
		}
	}
	
	protected static void _fftU01U3(float[][] src, float[] wi_list, float[] wj_list){
		int size = 0;
		int N = src.length;
		for (int x = src.length; x > 1; x = x >> 1){
			size++;
		}
		for (int q = 0; q < size; q++) {
			int step = 1 << q;
			int step2 = step << 1;
			for (int down = 0; down < N - step; down++) {
				for (int right = 0; right < N - step; right++) {
					boolean flag = (down/step) % 2 + (right/step) % 2 > 0 ? false : true;
					if (flag) {
						int wiInd = (down % step2)*(N/step2)*4;
						int wjInd = (right % step2)*(N/step2)*4;
						float[] wi = {wi_list[wiInd+1], wi_list[wiInd+3]};
						float[] wj = {wj_list[wjInd+1], wj_list[wjInd+3]};
						
						int ru1Ind = 4*right + 1;
						int ru3Ind = ru1Ind + 2;
						int rsu1Ind = 4*(right+step) + 1;
						int rsu3Ind = rsu1Ind+2;
						
						float a00u1 = src[down][ru1Ind];
						float a00u3 = src[down][ru3Ind];
						float a10u1 = src[down][ru1Ind];
						float a10u3 = src[down][ru3Ind];
						float a01u1 = src[down][ru1Ind];
						float a01u3 = src[down][ru3Ind];
						float a11u1 = src[down][ru1Ind];
						float a11u3 = src[down][ru3Ind];
						
						float[] t = new float[2];
						t[0] = 2*wi[0]*src[down+step][ru1Ind] - 2*wi[1]*src[down+step][ru3Ind];
						t[1] = 2*wi[0]*src[down+step][ru3Ind] + 2*wi[1]*src[down+step][ru1Ind];
						a00u1 += t[0];
						a00u3 += t[1];
						a10u1 -= t[0];
						a10u3 -= t[1];
						a01u1 += t[0];
						a01u3 += t[1];
						a11u1 -= t[0];
						a11u3 -= t[1];
										
						t[0] = 2*wj[0]*src[down][rsu1Ind] - 2*wj[1]*src[down][rsu3Ind];
						t[1] = 2*wj[0]*src[down][rsu3Ind] + 2*wj[1]*src[down][rsu1Ind];
						a00u1 += t[0];
						a00u3 += t[1];
						a10u1 += t[0];
						a10u3 += t[1];
						a01u1 -= t[0];
						a01u3 -= t[1];
						a11u1 -= t[0];
						a11u3 -= t[1];
						
						float[] t2 = new float[2];
						t[0] = 2*wj[0]*src[down+step][rsu1Ind] - 2*wj[1]*src[down+step][rsu3Ind];
						t[1] = 2*wj[0]*src[down+step][rsu3Ind] + 2*wj[1]*src[down+step][rsu1Ind];
						t2[0] = 2*wi[0]*t[0] - 2*wi[1]*t[1];
						t2[1] = 2*wi[0]*t[1] + 2*wi[1]*t[0];
						a00u1 += t2[0];
						a00u3 += t2[1];
						a10u1 -= t2[0];
						a10u3 -= t2[1];
						a01u1 -= t2[0];
						a01u3 -= t2[1];
						a11u1 += t2[0];
						a11u3 += t2[1];
						
						src[down][ru1Ind] = a00u1;
						src[down][ru3Ind] = a00u3;
						src[down+step][ru1Ind] = a10u1;
						src[down+step][ru3Ind] = a10u3;
						src[down][rsu1Ind] = a01u1;
						src[down][rsu3Ind] = a01u3;
						src[down+step][rsu1Ind] = a11u1;
						src[down+step][rsu3Ind] = a11u3;		
					}
				}
			}
		}
	}

	protected static void inverse(float[] arr){
		int dl = arr.length/2;
		int st = arr.length - 1;
		int j = 0;
		for (int i = 0; i < st; i++){
			if (i < j){
				float s0 = arr[i];
				arr[i] = arr[j];
				arr[j] = s0;
			}
			int k = dl;
			while (k <= j){
				j = j - k;
				k = k >> 1;
			}
			j = j + k;
		}
	}
	
	protected static void inverse(float[][] src){
		int N = src.length;
		for (int i = 0; i < N; i++){
			inverse(src[i]);
		}
		float[] arr = new float[N];
		for (int i = 0; i < N; i++){
			for (int j = 0; j < N; j++){
				arr[j] = src[j][i];
			}
			inverse(arr);
			for (int j = 0; j < N; j++){
				src[j][i] = arr[j];
			}
		}
	}
	
	protected static float[][] transformToB2Complex(float[][] src){
		float[][] res = new float[src.length][src.length*4];
		for (int i = 0; i < src.length; i++) {
			for (int j = 0; j < src.length; j++) {
				res[i][4*j] = src[i][j]/2;
				res[i][4*j+1] = src[i][j]/2;
				res[i][4*j+2] = 0;
				res[i][4*j+3] = 0;
			}
		}
		return res;
	}
	
	public static void main(String[] args) {
		float[][] arr = Generator.generateSimpleMatrix(8);
		Printer.printVector(arr);
		long time = System.nanoTime();
		System.out.println();
		inverse(arr);
		arr = transformToB2Complex(arr);
		_fft(arr);
		Printer.printVector(arr);
		time = System.nanoTime() - time;
		System.out.println("RESULT: " + time);
	}
	
}
